package org.hepeng.workx.spring.cloud.netflix.ribbon.client;

import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.client.loadbalancer.LoadBalancerRequestFactory;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;

import java.io.IOException;
import java.net.URI;
import java.util.Objects;

/**
 * @author he peng
 */
public class CompatibleLoadBalancerInterceptor implements ClientHttpRequestInterceptor {

    private LoadBalancerClient loadBalancer;
    private LoadBalancerRequestFactory requestFactory;

    public CompatibleLoadBalancerInterceptor(LoadBalancerClient loadBalancer) {
        this.loadBalancer = loadBalancer;
        this.requestFactory = new LoadBalancerRequestFactory(loadBalancer);
    }

    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
        final URI originalUri = request.getURI();
        String serviceName = originalUri.getHost();
        ServiceInstance serviceInstance = this.loadBalancer.choose(serviceName);
        ClientHttpResponse response;
        if (Objects.nonNull(serviceInstance)) {
            response = this.loadBalancer.execute(serviceName, requestFactory.createRequest(request, body, execution));
        } else {
            response = execution.execute(request , body);
        }
        return response;
    }
}
